Python: fix: keep AG-UI workflow reasoning in thread snapshots - #8058
Conversation
`_emit_text_reasoning` persists each workflow reasoning message into
`flow.reasoning_messages`, and `_WorkflowSnapshotBuilder.observe` folded
`TextMessage*` and `ToolCall*` events into the synthesized snapshot but had
no `Reasoning*` branch. Reasoning events fell through silently, so
intermediate output rendered live and then vanished when the thread was
hydrated from a snapshot -- while the same run through the agent path kept
it.
Fold reasoning into the builder: accumulate deltas per message_id and emit
entries in the shape the agent path already produces
({"id", "role": "reasoning", "content", ["encryptedValue"]}). Reasoning is
flushed at build() and at text-message and tool-call starts so a block that
streamed before other output replays in the position it streamed in.
Reasoning deliberately does not close an open tool-call group: it is UI-only
state that `agui_messages_to_agent_framework` drops, so it cannot break the
tool_calls/result adjacency providers require.
Fixes microsoft#8054
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GitHub Actions
Active
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
the
python
GitHub Actions
Active
There was a problem hiding this comment.
🟡 Changes recommended
The moderate encrypted-only reasoning loss must be fixed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Preserves AG-UI workflow reasoning in thread snapshots so hydrated output matches live streaming.
Changes:
- Folds reasoning events and encrypted values into workflow snapshots.
- Preserves reasoning output order.
- Adds unit and end-to-end regression coverage.
File summaries
| File | Summary |
|---|---|
python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py |
Tests reasoning persistence through snapshot hydration. |
python/packages/ag-ui/tests/ag_ui/test_snapshots.py |
Tests reasoning synthesis, ordering, and encryption. |
python/packages/ag-ui/agent_framework_ag_ui/_workflow.py |
Adds reasoning snapshot synthesis; encrypted-only reasoning can be dropped when the message-end event precedes its encrypted value. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
Address review: an encrypted-value-only reasoning message was dropped for the event order `_emit_text_reasoning` produces without a flow, where REASONING_MESSAGE_END precedes REASONING_ENCRYPTED_VALUE. The message carries no display text, so closing it at REASONING_MESSAGE_END discarded it and left the encrypted value with nothing to attach to. The underlying mistake was conflating two protocol levels. REASONING_MESSAGE_END closes the message; REASONING_END closes the block; and an encrypted value is block-scoped, so it legitimately trails the message end. Keep the message open past REASONING_MESSAGE_END and let REASONING_END, a new REASONING_START, intervening text/tool output, or build() finalize it. Handle REASONING_START so a new block also closes anything the previous one left open. Adds regression coverage for the reported order, for an empty block with neither text nor an encrypted value still being dropped, for a block closed by the next REASONING_START, for an encrypted value arriving after intervening text, and for end events naming a message that was never opened. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GitHub Actions
Active
| elif isinstance(event, ReasoningStartEvent): | ||
| # A new reasoning block supersedes anything still open from the last one. | ||
| self._flush_open_reasoning_message() | ||
| elif isinstance(event, ReasoningMessageStartEvent): | ||
| self._observe_reasoning_start(event) | ||
| elif isinstance(event, ReasoningMessageContentEvent): |
There was a problem hiding this comment.
What happens when a workflow emits an output event before a later intermediate or deprecated data event? The assistant text remains in _open_text_message, then build() flushes the newer reasoning first, so hydration reverses the order that streamed. Could either reasoning-start branch flush _open_text_message before opening the reasoning block?
There was a problem hiding this comment.
Confirmed, and reproduced as a failing test before changing anything: text streams, no
TextMessageEndEvent arrives because the workflow keeps running, reasoning opens, and build()
flushes reasoning first so hydration replays it before the text that streamed earlier.
Rather than reorder build(), I made the two slots peers so they can never both be open, which
removes the dependence on that order entirely. _observe_text_start and _observe_tool_call_start
already flushed open reasoning; reasoning-open paths now flush open text.
Working through it turned up three more instances of the same gap, all fixed here:
_observe_text_contentopens a message when text resumes without a start event, and did not flush
open reasoning. So in text -> reasoning -> text the reasoning still landed last._observe_tool_call_resultdid not flush open reasoning either, so reasoning that streamed before
a tool result replayed after it._observe_tool_call_startdoes flush, which is why this only
showed on the result path.- Same method, worse symptom:
_observe_text_contentalso opened a message over one already open
under a different id, discarding its content outright, where_observe_text_startflushes first.
Two content events with different ids and no start events keep only the second today. Concurrent
executors reach this by interleaving content events, which is exactly what the live run does, so I
fixed it in the same place rather than leave a silent text loss next to the flush I was adding.
I checked that the _observe_tool_call_result change is safe against the constraint the comment
there records. A reasoning block can now land between a tool call and its result, but agui_messages_to_agent_framework drops
role == "reasoning" before provider conversion, so the call and result come back adjacent. Nothing
asserted that across the two modules, so there is now a test that converts the snapshot and checks
the adjacency rather than trusting the comment.
Worth knowing that this ordering is not hypothetical: on a live fan-out workflow the reasoning, text
and tool events interleave across concurrently scheduled executors, and the alternation count varied
between 4 and 10 across three runs of the same prompt.
One consequence worth your call. Splitting an open text message means a message that resumes under
the same message_id would replay twice under that id. I re-identify the later fragment the way
_observe_tool_call_start already re-identifies a split message, and only when the id genuinely
collides, so a message that never resumes keeps the id it streamed under. That preserves streamed
order at the cost of one synthetic id.
The alternative is to merge the resumed fragment back into the earlier message, which keeps ids
stable but puts the reasoning after text that streamed later -- the defect this comment is about, in
a smaller form. I chose faithful order; say the word if you would rather have stable ids.
I checked the consumer side before settling on that, since splitting a message and changing an id
could plausibly break thread continuation. Two things make it safe, and both look deliberate rather
than accidental:
_snapshot_messages_matchonly requires equal ids for non-assistant roles. For an assistant
message with mismatched ids it falls through to_canonical_snapshot_message, which popsid
before comparing -- so assistant ids are already non-load-bearing for identity, which is what
makes the existing_observe_tool_call_startre-identification safe too.- A conforming client accumulates content per
message_idand so holds one merged message where
the snapshot now holds two fragments. Running_reconstruct_messages_from_thread_snapshotwith
that mismatch appends only the new user turn and does not duplicate the assistant history,
because reconstruction is backend-authoritative and client-supplied assistant messages are
filtered out of the incoming suffix.
So the split is contained to the stored shape and does not leak into the reconstructed transcript.
| def _observe_text_start(self, event: TextMessageStartEvent) -> None: | ||
| self._flush_open_reasoning_message() | ||
| if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id: |
There was a problem hiding this comment.
Following up on the earlier encrypted-only case: could we keep an empty ended reasoning message addressable when intervening text starts? This flush drops the empty shell, so a later ReasoningEncryptedValueEvent finds neither _open_reasoning_message nor a synthesized message and the protected value still disappears on hydration. The new late-value path handles intervening output only when the reasoning message also had visible text.
There was a problem hiding this comment.
You're right that the value is still lost, and this is a real hole in the earlier fix rather than an
edge case. Keeping the message open past REASONING_MESSAGE_END only helps when the value arrives
before anything else flushes it. A block carrying only protected data has no content event at all,
so intervening output flushed an empty message, and the late ReasoningEncryptedValueEvent then
found neither _open_reasoning_message nor a synthesized message to attach to. My earlier test only
covered the case where the reasoning also had visible text, which is exactly the gap you identified.
I went to a live fan-out workflow on gpt-5-mini to check the shape of this against a real provider,
and it refines your description in one way worth passing on. With a flow, the encrypted value is
emitted between REASONING_MESSAGE_START and REASONING_MESSAGE_END, not after the message end -- the
order in your comment is what the no-flow branch of _emit_text_reasoning produces. So on the flow
path the vulnerable window is narrower than stated, and it is a concurrency race rather than a
deterministic emission order.
It is still a real loss. The live run confirms both preconditions: reasoning arrives with an
encrypted value and no visible text at all, and concurrent executors interleave their events.
Replaying the captured event order with a concurrent text message starting in the gap before the
value arrives loses the encrypted reasoning on the current branch and keeps it with this change.
That is now a regression test built from the captured ordering rather than a synthesized one.
Two other things from that run: the provider emits ReasoningEncryptedValueEvent twice for one
reasoning item with the same entity_id, which is idempotent here but pinned by the test; and the
saved snapshot ends up with the reasoning row positioned where it streamed and dropped before
provider conversion, with no duplicate ids.
The fix itself: the flush now keeps the empty message addressable in the position it streamed, and
build() filters it out of the snapshot only while nothing has claimed it. Genuinely empty reasoning
is still absent, so the existing test that drops reasoning with neither text nor an encrypted value
still holds. I deliberately did not attach the value on arrival instead, which would have been the
smaller change: the message would then be appended after the output that flushed it, so the value
would survive but the reasoning would replay in the wrong place, trading this bug for the one in your
other comment. Position has to be decided at flush time.
The filter compares by object identity, not by id, so a caller-supplied message from raw_messages
that happens to share an id with a synthesized one cannot be removed with it. To be straight about
the strength of that choice: build() runs once per run at the single call site, so filtering rather
than deleting is not observable through the current path. I went that way because it keeps build()
a projection of accumulated state instead of mutating it -- design hygiene, not a second defect.
One consequential cleanup while in here: _observe_reasoning_start had an id check that my change
made dead, and following it through showed a repeated start for the block already open used to
discard the deltas already folded into it. It is now a no-op, with a test.
Evan Mattson (moonbox3)
added
the
ag-ui
Review on microsoft#8058 found two ordering defects in `_WorkflowSnapshotBuilder`. An `output` event followed by a later `intermediate` event left both the open text message and the open reasoning message unflushed, so `build()`'s fixed flush order decided the sequence and hydration reversed what streamed. Nothing flushed open text when reasoning opened, even though `_observe_text_start` and `_observe_tool_call_start` already flush open reasoning. Reasoning is now flushed wherever other output is appended and text is flushed wherever reasoning opens, so the two slots can never both be open and `build()`'s order cannot matter. That covers the `_observe_text_content` resume path and `_observe_tool_call_result`, which had the same gap. Reasoning is UI-only and dropped before provider conversion, so a block that now lands between a tool call and its result still converts back with the two adjacent, which provider APIs require; a test pins that across the two modules. A block carrying only protected data has no content event at all, so output arriving before its encrypted value flushed an empty message that the late `ReasoningEncryptedValueEvent` could no longer find, and the protected value was lost. The empty message now stays addressable in the position it streamed and `build()` filters it out only while nothing has claimed it, so genuinely empty reasoning is still absent from the snapshot. Filtering rather than deleting is not observable through the one call path today, which builds once per run; it keeps `build()` a projection of accumulated state so the choice does not have to be revisited if the builder is ever driven incrementally. A live fan-out workflow on gpt-5-mini shows both preconditions for that loss. Reasoning arrives with an encrypted value and no visible text, the value is emitted twice for one item, and with a flow it lands between REASONING_MESSAGE_START and REASONING_MESSAGE_END rather than after it. Concurrent executors interleave their events, so a text message can start in the gap before the value arrives; replaying the captured order with that interleaving loses the encrypted reasoning before this change and keeps it after. That same interleaving exposed a second loss on the no-start text path: `_observe_text_content` opened a message over one already open under a different id, discarding its content, where `_observe_text_start` flushes first. It now flushes as well. Splitting an open text message on reasoning meant a resumed message reusing its id replayed twice under that id; the later fragment is now re-identified the way `_observe_tool_call_start` already re-identifies a split message, and only when a collision is real. A repeated start for the block already open is now a no-op rather than silently discarding the deltas folded into it. Detecting that collision by scanning the accumulated messages on every flush made snapshot building quadratic -- flat at about 8us per message before, rising to 124us per message by 4000 messages. Appends now go through one method that maintains an id index, which restores linear scaling, and a test pins the index against a future append that bypasses it.
GitHub Actions
Active
GitHub Actions
Active
|
Two things that belong on the PR rather than in either review thread. A performance regression I introduced and fixed. Detecting the message-id collision by scanning Live validation on real Foundry infrastructure. Two of the questions in review are about what the
Replaying the captured event order with a concurrent text message in the gap before the value arrives Validation: |
Motivation & Context
With an AG-UI snapshot store enabled, a
Workflowrun's intermediate reasoning renders live and then disappears when the thread is hydrated. The same run through the agent path keeps it, so live and replayed output disagree.The reasoning is already recorded and simply never read back.
_emit_text_reasoning(_run_common.py:1051), which_emit_contentdelegates to fortext_reasoningcontent, persists each reasoning message intoflow.reasoning_messages— its docstring says so outright. The agent runner honours that contract (_agent_run.py:2100-2102, emitted terminally at:3202-3212); the workflow runner calls the same_emit_content, filling the same list, and never consumes it._WorkflowSnapshotBuilder.observefoldsTextMessage*andToolCall*events into the synthesized snapshot but had noReasoning*branch, so reasoning events fell through silently and never reached the store.Evan Mattson (@moonbox3) raised this in review on #8003 and confirmed a follow-up PR was the right home for it ("Follow up is fine, thanks."). This is that follow-up, built on top of #8003 now that it has merged.
Description & Review Guide
What are the major changes?
_WorkflowSnapshotBuildernow folds reasoning events, mirroring the existing text handling:observegains branches forReasoningMessageStartEvent,ReasoningMessageContentEvent,ReasoningMessageEndEvent/ReasoningEndEvent, andReasoningEncryptedValueEvent._observe_reasoning_*/_flush_open_reasoning_messagehelpers accumulate deltas permessage_idand emit entries in the same shape the agent path produces:{"id", "role": "reasoning", "content", ["encryptedValue"]}.Review found that the first version of this only ordered correctly for some event streams, so the
flush points are now symmetric rather than partial:
reasoning opens. The two slots can therefore never both be open, which is what makes
build()'sflush order stop deciding the sequence. Previously an
outputevent followed by a laterintermediateevent left both open and hydration replayed them in the wrong order._observe_text_content(text resuming with nostart event),
_observe_tool_call_result, and reasoning opened from a content event with nostart event.
its encrypted value used to flush an empty message that the late
ReasoningEncryptedValueEventcould no longer find, silently dropping the protected value. The empty message now stays
addressable in the position it streamed, and
build()leaves it out of the snapshot only whilenothing has claimed it, so genuinely empty reasoning is still absent.
message_idwould replaytwice under that id. The later fragment is re-identified the way
_observe_tool_call_startalready re-identifies a split message, and only when a collision is real.
_observe_text_contentopened a message over onealready open under a different id and discarded its content, where
_observe_text_startflushesfirst. It now flushes as well.
What is the impact of these changes?
Reasoning that streamed during a workflow run is still present after the thread is hydrated from a
snapshot, and it replays in the position it streamed rather than wherever
build()happened to putit.
Reasoning still does not close an open tool-call group -- verified, two tool calls with a reasoning
block between them stay in one assistant message. A reasoning row can now sit between a tool call
and its result, which is safe because
_message_adapters.py:691-694dropsrole == "reasoning"when converting back to provider messages, so the call and result come back adjacent. Nothing
asserted that across the two modules before; a test now converts the snapshot and checks it.
One behaviour change worth calling out: a text message that resumes after an interleaved reasoning
block replays as two messages, and the later fragment carries a generated id rather than the one it
streamed under. A message that never resumes keeps its original id.
This is the smaller of the two shapes discussed on the issue. It reads the real event stream rather
than
flow, which matters for workflows: request_info/interrupt tool calls and executorpassthrough events are yielded directly and never touch
flow, soobservesees strictly morethan
_build_messages_snapshotwould. The second shape -- emitting a terminalMessagesSnapshotEventfrom the workflow runner -- is left to the issue, since an emitted snapshotis treated as authoritative and would need to reproduce those bypassing events too.
What do you want reviewers to focus on?
Two judgment calls rather than the mechanics.
First, the re-identification above. It preserves streamed order at the cost of one generated id.
The alternative is to merge the resumed fragment back into the earlier message, which keeps ids
stable but replays the reasoning after text that streamed later -- a smaller version of the defect
this PR fixes. Say if you would rather have stable ids.
Second, the empty-reasoning message being filtered at
build()rather than dropped at flush.Filtering is not observable through the single call path today, which builds once per run; it keeps
build()a projection of accumulated state rather than a mutation of it. Attaching the value onarrival instead would have been smaller but would replay the reasoning after the output that
flushed it.
A note on measurement, since two of these depend on what the provider and scheduler actually do
rather than on the folding logic: a live fan-out workflow on
gpt-5-minishows reasoning arrivingwith an encrypted value and no visible text, the same value emitted twice for one item, and
concurrent executors interleaving their events. Details are in a PR comment.
Related Issue
Fixes #8054
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.